package 平衡二叉树;

public class AVLNode {
    private int data;//数据
    private int depth; //保存节点高度
    private int balance; //是否平衡
    public AVLNode parent; //父节点
    public AVLNode left; //左子节点
    public AVLNode right; //右子节点

    public AVLNode(int data){
        this.data=data;
        depth=1;
        left=null;
        right=null;
    }


    //插入数据
    public void insert(AVLNode root,int data){
        //小于根节点 插入左边
        if(data<root.data){
            if(root.left!=null){
                insert(root.left,data);
            }else{
                root.left=new AVLNode(data);
                root.left.parent=root;
            }
        }
        //大于根节点 插入右边
        else{
            if (root.right!=null){
                insert(root.right,data);
            }else{
                root.right=new AVLNode(data);
                root.right.parent=root;
            }
        }
        //插入之后计算平衡度
        root.balance=calcBalance(root);
        //左子树高
        if(root.balance>=2){
            //如果左子树是-1 ，则先左旋后右旋，否则直接右旋
            if(root.left.balance==-1){
                left_rotate(root.left);
            }
            right_rotate(root);
        }
        //右子树高
        if(root.balance<=-2){
            //如果右子树是1，则先右旋再左旋，否则直接左旋
            if(root.left.balance==1){
                right_rotate(root.right);
            }
            left_rotate(root);
        }
        //调整之后，重新计算平衡因子
        root.balance=calcBalance(root);
        root.depth=calcDepth(root);
    }

    private int calcDepth(AVLNode p) {
        int depth=0;
        if(p.left!=null){
            depth=p.left.depth;
        }
        if(p.right!=null && depth<p.right.depth){
            depth=p.right.depth;
        }
        depth++;
        return depth;
    }

    private int calcBalance(AVLNode p) {
        int left_depth;
        int right_depth;
        //左子树深度
        if(p.left!=null){
            left_depth=p.left.depth;
        }else {
            left_depth=0;
        }
        //右子树深度
        if (p.right!=null){
            right_depth=p.right.depth;
        }else {
            right_depth=0;
        }

        return left_depth-right_depth;
    }

    //右旋
    private void right_rotate(AVLNode p) {
        AVLNode pParent=p.parent;
        AVLNode pLeftSon=p.left;
        AVLNode pRightGrandSon=pLeftSon.right;
        //左子变父
        pLeftSon.parent=pParent;
        if(pParent!=null){
            if(p==pParent.left){
                pParent.left=pLeftSon;
            }else if (p==pParent.right){
                pParent.right=pLeftSon;
            }
        }
        pLeftSon.right=p;
        p.parent=pLeftSon;
        p.left=pRightGrandSon;
        if (pRightGrandSon!=null){
            pRightGrandSon.parent=p;
        }
        p.depth=calcDepth(p);
        p.balance=calcBalance(p);
        pLeftSon.depth=calcDepth(pLeftSon);
        pLeftSon.balance=calcBalance(pLeftSon);
    }
    //左旋
    private void left_rotate(AVLNode p) {
        AVLNode pParent=p.parent;
        AVLNode pRightSon=p.right;
        AVLNode pLeftGrandSon=pRightSon.left;
        //左子变父
        pRightSon.parent=pParent;
        if(pParent!=null){
            if(p==pParent.right){
                pParent.right=pRightSon;
            }else if (p==pParent.left){
                pParent.left=pRightSon;
            }
        }
        pRightSon.left=p;
        p.parent=pRightSon;
        p.right=pLeftGrandSon;
        if (pLeftGrandSon!=null){
            pLeftGrandSon.parent=p;
        }
        p.depth=calcDepth(p);
        p.balance=calcBalance(p);
        pRightSon.depth=calcDepth(pRightSon);
        pRightSon.balance=calcBalance(pRightSon);
    }
}
